"""李华自定义的"""

from __future__ import annotations

import datetime
import time
from datetime import timedelta
from typing import List
from typing import Optional
from typing import Union

import numpy as np
import pandas as pd
import pandas_market_calendars as tc

import os
import sys
ROOT = os.path.dirname(os.path.abspath("."))
sys.path.append(ROOT)
import os.path as osp
# 尝试导入akshare，如果不存在则提示安装
try:
    import akshare as ak
except ImportError:
    raise ImportError(
        "请安装akshare库: pip install akshare 或 conda install -c conda-forge akshare"
    )

ALL_INDICATORS = [
    "macd",         # MACD指标
    "boll_ub",      # 布林带上轨
    "boll_lb",      # 布林带下轨
    "rsi_30",       # 30日相对强弱指数
    "cci_30",       # 30日商品通道指数
    "dx_30",        # 30日方向运动指数
    "ma_price_7",
    "ma_price_13",
    "ma_price_30",
    "ma_price_60",
    "ma_price_90",
    "ma_price_180",
    "ma_price_300",

    "ma_price_7",
    "ma_price_13",
    "ma_price_30",
    "ma_price_60",
    "ma_price_90",
    "ma_price_180",
    "ma_price_300",

    "slope_price_7",
    "slope_price_13",
    "slope_price_30",
    "slope_price_60",
    "slope_price_90",
    "slope_price_180",
    "slope_price_300",

    "ma_volume_7",
    "ma_volume_13",
    "ma_volume_30",
    "ma_volume_60",
    "ma_volume_90",
    "ma_volume_180",
    "ma_volume_300",

    "slope_volume_7",
    "slope_volume_13",
    "slope_volume_30",
    "slope_volume_60",
    "slope_volume_90",
    "slope_volume_180",
    "slope_volume_300",
]

# 定义一个0.0数组 长度和ALL_INDICATORS一样长
zero_values = np.zeros(len(ALL_INDICATORS)).tolist()
class LihuaProcessor:
    """Provides methods for retrieving daily stock data from
    Akshare API for Chinese A-share market
    """

    def __init__(self):
        """
        李华实现的从文件读取
        """
        pass


    def download_data(
        self,
        ticker_list: List[str],
        start_date: str,
        end_date: str,
        time_interval: str,
        proxy: Optional[str] = None,
        adj: str = "",
    ) -> pd.DataFrame:

        # 设置参数
        self.start = start_date
        self.end = end_date
        self.time_interval = time_interval


        # 下载并保存数据到pandas DataFrame
        start_date_dt = pd.Timestamp(start_date)
        end_date_dt = pd.Timestamp(end_date)
        delta = timedelta(days=1)
        data_df = pd.DataFrame()

        # 将列转换为合适的数据类型
        def safe_convert(value, dtype):
            try:
                if isinstance(value, str):
                    return dtype(value.replace(",", ""))
                return value
            except :
                return value

        for tic in ticker_list:
            current_tic_start_date = start_date_dt

            # 日级数据 后期有其他级别的数据 修改数据目录
            data_dir = osp.join(ROOT, 'data', 'stock_daily_data')
            try:
                print(f"处理股票 {tic}")
                file_path = os.path.join(data_dir, f"{tic.replace('.', '_')}.csv")
                temp_df = pd.read_csv(file_path)
                temp_df['timestamp'] = pd.to_datetime(temp_df['date'])
                # 过滤timestamp 不在start_date_dt 和 end_date_dt之间的数据
                temp_df = temp_df[(temp_df['timestamp'] >= start_date_dt) & (temp_df['timestamp'] <= end_date_dt)]

                # 获取除tic列之外的列名
                other_cols = [col for col in temp_df.columns if col != 'tic']
                # 重新指定列顺序
                temp_df = temp_df[other_cols + ['tic']]

                temp_df["open"] = temp_df["open"].apply(lambda x: safe_convert(x, float))
                temp_df["high"] = temp_df["high"].apply(lambda x: safe_convert(x, float))
                temp_df["low"] = temp_df["low"].apply(lambda x: safe_convert(x, float))
                temp_df["close"] = temp_df["close"].apply(lambda x: safe_convert(x, float))

                temp_df["volume"] = temp_df["volume"].apply(lambda x: safe_convert(x, int))

                temp_df["amount"] = temp_df["amount"].apply(lambda x: safe_convert(x, float))
                temp_df["amplitude"] = temp_df["amplitude"].apply(lambda x: safe_convert(x, float))
                temp_df["change_rate"] = temp_df["change_rate"].apply(lambda x: safe_convert(x, float))
                temp_df["change_amount"] = temp_df["change_amount"].apply(lambda x: safe_convert(x, float))
                temp_df["turnover_rate"] = temp_df["turnover_rate"].apply(lambda x: safe_convert(x, float))
                ma_periods = [7, 13, 30, 60, 90, 180, 300]
                for period in ma_periods:

                    temp_df[f'ma_price_{period}'] = temp_df[f'ma_price_{period}'].apply(lambda x: safe_convert(x, float))
                    temp_df[f'slope_price_{period}'] = temp_df[f'slope_price_{period}'].apply(lambda x: safe_convert(x, float))
                    temp_df[f'ma_volume_{period}'] = temp_df[f'ma_volume_{period}'].apply(lambda x: safe_convert(x, float))
                    temp_df[f'slope_volume_{period}'] = temp_df[f'slope_volume_{period}'].apply(lambda x: safe_convert(x, float))

                temp_df["macd"] = temp_df["macd"].apply(lambda x: safe_convert(x, float))
                temp_df["boll_ub"] = temp_df["boll_ub"].apply(lambda x: safe_convert(x, float))
                temp_df["boll_lb"] = temp_df["boll_ub"].apply(lambda x: safe_convert(x, float))
                temp_df["rsi_30"] = temp_df["rsi_30"].apply(lambda x: safe_convert(x, float))
                temp_df["cci_30"] = temp_df["cci_30"].apply(lambda x: safe_convert(x, float))
                temp_df["dx_30"] = temp_df["dx_30"].apply(lambda x: safe_convert(x, float))

                # 添加股票代码列

                # 合并数据
                data_df = pd.concat([data_df, temp_df], ignore_index=True)


            except Exception as e:
                print(f"获取股票 {tic} 数据时出错: {e}")

        # 重命名列以匹配Yahoo Finance格式
        if  data_df.empty:
            # 如果没有获取到数据，创建一个空的DataFrame
            data_df = pd.DataFrame(
                columns=["timestamp", "open", "high", "low", "close", "volume", "tic"]
            )
        return data_df


    def clean_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        清理数据，处理缺失值和异常值

        参数:
        df : pd.DataFrame
            原始数据DataFrame

        返回:
        pd.DataFrame
            清理后的数据DataFrame
        """
        # 确保列名正确
        tic_list = np.unique(df.tic.values)
        df = df.copy()
        df = df.rename(columns={col: col.lower() for col in df.columns})
        times = self.get_trading_days(start=self.start, end=self.end)
        # 确保必要的列存在
        required_columns = ["timestamp", "open", "high", "low", "close", "volume", "tic"]
        for col in required_columns:
            if col not in df.columns:
                raise ValueError(f"缺少必要的列: {col}")


        # 按时间和股票代码排序
        df = df.sort_values(by=["tic", "timestamp"])

        # 创建一个包含完整时间戳序列的新 DataFrame
        new_df = pd.DataFrame()
        for tic in tic_list:
            tmp_df = pd.DataFrame(
                columns=["open", "high", "low", "close", "volume"]+ALL_INDICATORS, index=times
            )
            tic_df = df[
                df.tic == tic
            ]  # 从下载的数据中提取与该股票代码相关的行
            for i in range(tic_df.shape[0]):  # 使用原始数据填充空 DataFrame
                tmp_timestamp = tic_df.iloc[i]["timestamp"]
                tmp_df.loc[tmp_timestamp.strftime("%Y-%m-%d")] = tic_df.iloc[i][
                    ["open", "high", "low", "close", "volume"]+ALL_INDICATORS
                ]
                a=1+1

            if str(tmp_df.iloc[0]["close"]) == "nan":
                print("起始日期数据为 NaN，使用第一个有效数据填充。")
                for i in range(tmp_df.shape[0]):
                    if str(tmp_df.iloc[i]["close"]) != "nan":
                        first_valid_close = tmp_df.iloc[i]["close"]
                        tmp_df.iloc[0] = [
                            first_valid_close,
                            first_valid_close,
                            first_valid_close,
                            first_valid_close,
                            0.0,
                        ]+zero_values
                        break

            # 如果第一行的收盘价仍然为 NaN（这种情况下所有价格均为 NaN）
            if str(tmp_df.iloc[0]["close"]) == "nan":
                print(
                    "股票代码: ",
                    tic,
                    " 数据缺失。所有价格均为 NaN，使用 0 填充。",
                )
                tmp_df.iloc[0] = [
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                ]+zero_values
            # 用前一个收盘价填充 NaN 数据，并将成交量设为 0
            for i in range(tmp_df.shape[0]):
                if str(tmp_df.iloc[i]["close"]) == "nan":
                    previous_close = tmp_df.iloc[i - 1]["close"]
                    if str(previous_close) == "nan":
                        raise ValueError
                    tmp_df.iloc[i] = [
                        previous_close,
                        previous_close,
                        previous_close,
                        previous_close,
                        0.0,
                    ]+zero_values

            tmp_df = tmp_df.astype(float)
            tmp_df["tic"] = tic
            new_df = pd.concat([new_df, tmp_df])

        # 重置索引

        new_df = new_df.reset_index()
        new_df = new_df.rename(columns={"index": "timestamp"})

        return new_df

    def add_technical_indicator(
        self, data: pd.DataFrame, tech_indicator_list: List[str]
    ) -> pd.DataFrame:
        df = data.copy()
        df = df.sort_values(by=["tic", "timestamp"])
        return df[ ["timestamp", "open", "high", "low", "close", "volume", "tic"]+tech_indicator_list]

    def add_vix(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        添加波动率指数(VIX)
        注意：Akshare中没有直接的VIX数据，这里提供一个空实现以保持接口一致性

        参数:
        data : pd.DataFrame
            股票数据DataFrame

        返回:
        pd.DataFrame
            添加VIX列后的DataFrame
        """
        df = data.copy()
        print("警告: Akshare数据源不支持VIX数据，将添加默认值0")
        df["VIXY"] = 0
        return df

    def calculate_turbulence(
        self, data: pd.DataFrame, time_period: int = 252
    ) -> pd.DataFrame:
        """
        有bug有空在解决
        计算湍流指数

        参数:
        data : pd.DataFrame
            股票数据DataFrame
        time_period : int
            计算湍流指数的时间窗口，默认为252个交易日

        返回:
        pd.DataFrame
            包含时间戳和湍流指数的DataFrame
        """
        # 复制输入数据，避免修改原始数据
        df = data.copy()
        # 将数据转换为透视表格式，行索引为时间戳，列索引为股票代码，值为收盘价
        df_price_pivot = df.pivot(index="timestamp", columns="tic", values="close")
        # 计算每日收益率，用于后续的湍流指数计算
        df_price_pivot = df_price_pivot.pct_change()
        df_price_pivot.fillna(0, inplace=True);
        # 获取数据中的所有唯一日期
        unique_date = df.timestamp.unique()
        # 设置计算起始点，从第time_period个时间点开始，确保有足够的历史数据
        start = time_period
        # 初始化湍流指数列表，前start个时间点的湍流指数设为0
        turbulence_index = [0] * start
        # 用于记录连续有效湍流值的计数器
        count = 0
        # 从start开始遍历所有日期
        for i in range(start, len(unique_date)):
            # 获取当前日期的收益率数据
            current_price = df_price_pivot[df_price_pivot.index == unique_date[i]]
            # 获取当前日期前time_period个交易日的历史收益率数据
            hist_price = df_price_pivot[
                (df_price_pivot.index < unique_date[i])
                & (df_price_pivot.index >= unique_date[i - time_period])
            ]
            # 删除缺失值超过"最老"股票数量的股票，即从缺失值最少的行开始，删除有缺失值的列
            filtered_hist_price = hist_price.iloc[
                hist_price.isna().sum().min() :
            ].dropna(axis=1)

            # 计算过滤后历史数据的协方差矩阵
            cov_temp = filtered_hist_price.cov()
            # 计算当前收益率相对于历史收益率均值的偏差
            current_temp = current_price[[x for x in filtered_hist_price]] - np.mean(
                filtered_hist_price, axis=0
            )

            # 添加异常处理来解决SVD不收敛问题
            try:
                # 计算马氏距离，作为湍流指数的基础值
                temp = current_temp.values.dot(np.linalg.pinv(cov_temp)).dot(
                    current_temp.values.T
                )
            except np.linalg.LinAlgError:
                # 当SVD不收敛时，使用替代方法计算，输出警告信息并将湍流值设为0
                print(f"警告: 在时间点 {unique_date[i]} 的湍流指数计算中遇到SVD不收敛问题，使用默认值0")
                temp = np.array([[0]])
            except Exception as e:
                # 处理其他可能的异常，输出警告信息并将湍流值设为0
                print(f"警告: 在时间点 {unique_date[i]} 的湍流指数计算中遇到异常: {str(e)}，使用默认值0")
                temp = np.array([[0]])

            # 如果计算得到的湍流值大于0
            if temp > 0:
                # 计数器加1
                count += 1
                # 如果连续有效湍流值的数量大于2
                if count > 2:
                    # 使用计算得到的湍流值
                    turbulence_temp = temp[0][0]
                else:
                    # 避免因计算刚开始而出现大的异常值，将湍流值设为0
                    turbulence_temp = 0
            else:
                # 如果计算得到的湍流值小于等于0，将湍流值设为0
                turbulence_temp = 0
            # 将当前日期的湍流值添加到湍流指数列表中
            turbulence_index.append(turbulence_temp)

        # 将湍流指数列表转换为DataFrame，包含时间戳和湍流指数两列
        turbulence_index = pd.DataFrame(
            {"timestamp": df_price_pivot.index, "turbulence": turbulence_index}
        )
        return turbulence_index

    def add_turbulence(
        self, data: pd.DataFrame, time_period: int = 252
    ) -> pd.DataFrame:
        """
        添加从预计算的DataFrame中获取的湍流指数

        参数:
        data : pd.DataFrame
            股票数据DataFrame
        time_period : int
            计算湍流指数的时间窗口

        返回:
        pd.DataFrame
            添加湍流指数后的DataFrame
        """
        #data["turbulence"]=0
        df = data.copy()
        turbulence_index = self.calculate_turbulence(df, time_period=time_period)
        df = df.merge(turbulence_index, on="timestamp")
        df = df.sort_values(["timestamp", "tic"]).reset_index(drop=True)
        return df

    def df_to_array(
        self, df: pd.DataFrame, tech_indicator_list: List[str], if_vix: bool
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        将DataFrame转换为数组

        参数:
        df : pd.DataFrame
            股票数据DataFrame
        tech_indicator_list : list
            技术指标列表
        if_vix : bool
            是否包含VIX数据

        返回:
        tuple: (价格数组, 技术指标数组, 湍流数组)
        """
        price_array = []
        tech_array = []
        turbulence_array = []
        df = df.copy()
        # 获取所有唯一的股票代码
        unique_ticker = df.tic.unique()
        # 标志变量，用于判断是否是第一次循环
        if_first_time = True
        # 遍历每一只股票
        for tic in unique_ticker:
            # 第一次循环时，初始化各个数组
            # 这个if-else结构的目的是：第一次循环时初始化数组，后续循环时将数据横向拼接

            if if_first_time:
                # 第一只股票的价格数据，只包含收盘价列
                price_array = df[df.tic == tic][["close"]].values
                # 第一只股票的技术指标数据
                tech_array = df[df.tic == tic][tech_indicator_list].values
                # 根据if_vix参数决定使用哪种湍流数据
                if if_vix:
                    # 如果if_vix为True，使用VIXY列（波动率数据）
                    turbulence_array = df[df.tic == tic]["VIXY"].values
                else:
                    # 如果if_vix为False，使用turbulence列（湍流数据）
                    turbulence_array = df[df.tic == tic]["turbulence"].values
                # 将标志设为False，表示第一次循环已完成
                if_first_time = False
            else:
                # 获取当前股票的数据
                current_close = df[df.tic == tic][["close"]].values
                current_tech = df[df.tic == tic][tech_indicator_list].values

                # 获取当前数组和新数据的长度
                current_length = price_array.shape[0]
                new_length = current_close.shape[0]

                if new_length < current_length:
                    # 如果新数据比当前数组短，在新数据后面补NaN
                    pad_length = current_length - new_length
                    current_close = np.vstack([current_close, np.full((pad_length, 1), np.nan)])
                    current_tech = np.vstack([current_tech, np.full((pad_length, len(tech_indicator_list)), np.nan)])
                elif new_length > current_length:
                    # 如果新数据比当前数组长，截断新数据
                    current_close = current_close[:current_length]
                    current_tech = current_tech[:current_length]

                # 对于非第一只股票，将数据横向拼接（按列拼接）
                # 将当前股票的价格数据拼接到已有的price_array右侧
                price_array = np.hstack([price_array, current_close])
                # 将当前股票的技术指标数据拼接到已有的tech_array右侧
                tech_array = np.hstack([tech_array, current_tech])
                # 注意：这里只拼接了价格和技术指标数据，但没有拼接湍流数据
                # 因为在强化学习环境中，通常只需要一个全局的市场湍流值，而不是每只股票的湍流值
        # 返回处理后的三个数组
        return price_array, tech_array, turbulence_array

    def get_trading_days(self, start: str, end: str) -> List[str]:
        """
        获取交易日列表
        注意：对于中国A股市场，需要使用中国交易日历

        参数:
        start : str
            开始日期
        end : str
            结束日期

        返回:
        list: 交易日列表
        """
        # 获取中国A股交易日历
        try:
            csi = tc.get_calendar("SSE")  # 上交所交易日历
            df = csi.date_range_htf("1D", pd.Timestamp(start), pd.Timestamp(end))
            trading_days = []
            for day in df:
                trading_days.append(str(day)[:10])
            return trading_days
        except Exception as e:
            print(f"获取交易日历失败: {e}，将使用默认日期范围")
            # 如果无法获取交易日历，使用简单的日期范围
            start_date = pd.Timestamp(start)
            end_date = pd.Timestamp(end)
            date_range = pd.date_range(start=start_date, end=end_date, freq='D')
            return [str(date.date()) for date in date_range]

    def fetch_latest_data(
        self,
        ticker_list: List[str],
        time_interval: str,
        tech_indicator_list: List[str],
        limit: int = 100,
        adj: str = "",
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        获取最新的股票数据

        参数:
        ticker_list : list
            股票代码列表
        time_interval : str
            时间间隔
        tech_indicator_list : list
            技术指标列表
        limit : int
            获取的数据点数量限制
        adj : str
            复权方式

        返回:
        tuple: (最新价格, 最新技术指标, 最新湍流值)
        """
        time_interval = self.convert_interval(time_interval)

        end_datetime = datetime.datetime.now()
        start_datetime = end_datetime - datetime.timedelta(days=limit)

        data_df = pd.DataFrame()
        for tic in ticker_list:
            nonstandard_tic = self.transfer_standard_ticker_to_nonstandard(tic)
            try:
                # 使用Akshare获取近期数据
                barset = ak.stock_zh_a_hist(
                    symbol=nonstandard_tic,
                    period=time_interval,
                    start_date=start_datetime.strftime('%Y%m%d'),
                    end_date=end_datetime.strftime('%Y%m%d'),
                    adjust=adj,
                )
                barset["tic"] = tic
                data_df = pd.concat([data_df, barset], ignore_index=True)
                time.sleep(0.1)  # 避免请求过于频繁
            except Exception as e:
                print(f"获取股票 {tic} 最新数据时出错: {e}")

        if data_df.empty:
            raise ValueError("未能获取到任何最新数据")

        # 重命名列以匹配标准格式
        data_df.columns = [
            "timestamp",
            "open",
            "close",
            "high",
            "low",
            "volume",
            "amount",
            "amplitude",
            "pct_chg",
            "change",
            "turnover",
            "tic",
        ]
        data_df = data_df[
            ["timestamp", "open", "high", "low", "close", "volume", "tic"]
        ]
        data_df["timestamp"] = pd.to_datetime(data_df["timestamp"])

        # 按时间戳排序
        data_df = data_df.sort_values("timestamp").reset_index(drop=True)

        # 添加技术指标
        df = self.add_technical_indicator(data_df, tech_indicator_list)
        df["VIXY"] = 0  # Akshare不支持VIX数据

        # 转换为数组
        price_array, tech_array, turbulence_array = self.df_to_array(
            df, tech_indicator_list, if_vix=True
        )
        latest_price = price_array[-1]
        latest_tech = tech_array[-1]
        latest_turb = turbulence_array[-1]
        return latest_price, latest_tech, latest_turb